package net.tp.struct.bst;

import net.tp.struct.tree.TreePrinter;
import net.tp.struct.tree.BinNode;
import static net.tp.util.Utils.*;

import java.util.Arrays;
import java.util.Comparator;
import java.util.Random;

/**
 *
 * <p>Randomized Binary Search Tree.</p>
 *
 * <p>Don't be confused between random BST and randomized BST.</p>
 *
 * <p>Random BST is built using random insertions. Random BST has expected logarithmic height, hence has expected $O(log(n))$
 * time for search and update operations. However, the input is not always random. Randomized BST is an approach to cope
 * with this. Other randomized approaches include Treap and SkipList.</p>
 *
 * <p>If a tree is random, any key has a probability $1/n$ of being the root. So when inserting a new key, the new key
 * should have some chance of being the root of the tree, or the root of a subtree. The probability
 * of this node being at the root is $1/{size+1}$ where $size$ is the size of the tree or subtree. So for this data
 * structure to work, we need to maintain the size of the subtree in each node.</p>
 *
 * <p>Randomized BST have the same distribution as BST under random insertion no matter in what order keys are inserted.
 * To implement randomized BST, we need to maintain the subtree size at each node.</p>
 *
 * <h3>Insertion operation</h3>
 *
 * <p>Let's insert key D into this tree:</p>
 *
 * <pre>
 *      E:5
 *    ┌──┴──┐
 *   B:3   G:1
 *  ┌─┴─┐
 * A:1 C:1
 * </pre>
 *
 * <p>If this is a normal BST, D will end up in a leaf:</p>
 *
 * <pre>
 *         E:6
 *     ┌────┴────┐
 *    B:4       G:1
 *  ┌──┴──┐
 * A:1   C:2
 *      ┌─┴─┐
 *         D:1
 * </pre>
 *
 * <p>In the search path E-B-C, D has 1/6 chance of being root of tree E; if not, it has 1/4 chance of being root of subtree
 * B; if not, it has 1/2 chance of being root of subtree C, if not D will end up in the leaf node as above:</p>
 *
 * <pre>
 *         E:6 ← 1/6 chance
 *     ┌────┴─────────────┐
 *    B:4 ← 1/4 chance   G:1
 *  ┌──┴──┐
 * A:1   C:2 ← 1/2 chance
 *      ┌─┴─┐
 *         D:1
 * </pre>
 *
 * <p>If D wins the chance of being the root of tree E, D will be inserted into BST as above, then rotated up to the root.</p>
 *
 * <pre>
 *       D:6
 *    ┌───┴───┐
 *   B:3     E:2
 *  ┌─┴─┐   ┌─┴─┐
 * A:1 C:1     G:1
 * </pre>
 *
 * <p>If D does not win the chance of being the root of tree E, and wins the chance of being the root of subtree
 * B, D will be inserted in BST as above, then rotated up to the root of subtree B.</p>
 *
 * <pre>
 *           E:6
 *       ┌────┴────┐
 *      D:4       G:1
 *    ┌──┴──┐
 *   B:3
 *  ┌─┴─┐
 * A:1 C:1
 * </pre>
 *
 * <p>If D does not win the chance of being the root of tree E and subtree B, and wins the chance of being the root
 * of the subtree C, D is inserted as above, then rotated with C to be the root of the subtree.</p>
 *
 * <pre>
 *         E:6
 *     ┌────┴────┐
 *    B:4       G:1
 *  ┌──┴──┐
 * A:1   D:2
 *      ┌─┴─┐
 *     C:1
 * </pre>
 *
 * <h3>Deletion operation</h3>
 *
 * <p>For a normal BST, it's trivial to delete the node with 0 or 1 child, to delete an internal node with 2 children,
 * we first need to find the successor of the node (which must be either a leaf node or a node with 1 child),
 * swap the elements between the node and the successor node, then delete the node holding the successor value.</p>
 *
 * <p>For the randomized BST, the same deletion operation is for node with 0 or 1 child. For the node with 2 children,
 * randomized BST uses a join strategy based on probability.</p>
 *
 * <p>To delete node D of this tree:</p>
 *
 * <pre>
 *       D:6
 *    ┌───┴───┐
 *   B:3     E:2
 *  ┌─┴─┐   ┌─┴─┐
 * A:1 C:1     G:1
 * </pre>
 *
 * <p>Deleting node D will turn the tree into 2 subtrees that need to be joined:</p>
 *
 * <pre>
 *   B:3 ← 3/5 chance          E:2 ← 2/5 chance of being root
 *  ┌─┴─┐                     ┌─┴─┐
 * A:1 C:1                       G:1
 * </pre>
 *
 * <p>Assuming B win the chance of being the root, we need to join the 2 subtree rooted at C and E:</p>
 *
 * <pre>
 *   B:?
 *  ┌─┴
 * A:1       C:1 ← 1/3 chance     E:2 ← 2/3 chance of being the root of the subtree under B
 *                               ┌─┴─┐
 *                                  G:1
 * </pre>
 *
 * <p>If E wins the chance of being the root of the subtree:</p>
 *
 * <pre>
 *     B:5
 *  ┌───┴───┐
 * A:1     E:3
 *        ┌─┴─┐
 *       C:1 G:1
 * </pre>
 *
 * <p>If C wins the chance of being the root of the subtree:</p>
 *
 * <pre>
 *     B:5
 *  ┌───┴───┐
 * A:1     C:3
 *        ┌─┴─┐
 *           E:2
 *          ┌─┴─┐
 *             G:1
 * </pre>
 *
 *
 *
 *
 *
 *
 * @author Trung Phan
 * @see <p>Conrado Martinez and Salvador Roura, Randomized Binary Search Trees, Journal of the ACM, 45(2):288-323, March 1998</p>
 *
 */
public class RandomizedBST<E> {


    private final Comparator<? super E> comparator;

    private BinNode<E> root;
    private Random random = new Random();

    public RandomizedBST() {
        this(naturalComparator());
    }

    public RandomizedBST(Comparator<? super E> comparator) {
        this.comparator = comparator;
    }

    public void add(E element) {

        if (root == null) {
            root = new BinNode<>(element);
        }
        else {
            BinNode<E> n = root;
            while (true) {
                if (random.nextInt(n.size + 1) == 0) {
                    addRoot(n, element);
                    break;
                }

                int cmp = comparator.compare(element, n.element);
                if (cmp == 0) {
                    // element exists : ignore
                    break;
                }
                BinNode<E> next = cmp < 0 ? n.left : n.right;
                if (next != null) n = next;
                else {
                    if (cmp < 0) n.bindLeft(new BinNode<>(element));
                    else n.bindRight(new BinNode<>(element));

                    for (BinNode<E> i = n; i != null; i.size++, i = i.parent);
                    break;
                }
            }
        }

    }

    private void addRoot(BinNode<E> node, E element) {
        BinNode<E> parent = node.parent;

        BinNode<E> n = find(node, element);
        int cmp = comparator.compare(element, n.element);
        if (cmp != 0) {
            if (cmp < 0) n.bindLeft(new BinNode<>(element));
            else n.bindRight(new BinNode<>(element));

            for (BinNode<E> i = n; i != null; i.size++, i = i.parent);

            n = cmp < 0 ? n.left : n.right;
        }

        while (n.parent != parent) {
            if (n.parent.left == n) n.parent.rotateRight();
            else n.parent.rotateLeft();

        }

        if (n.parent == null) root = n;
    }


    public E remove(E element) {
        BinNode<E> n = find(root, element);
        if (n == null || comparator.compare(element, n.element) != 0) return null;
        BinNode<E> parent = n.parent;

        for (BinNode<E> i = parent; i != null; i.size--, i = i.parent);

        BinNode<E> join = join(n.left, n.right);

        if (parent == null) root = join;
        else if (parent.left == n) parent.bindLeft(join);
        else parent.bindRight(join);

        return n.element;
    }

    private BinNode<E> join(BinNode<E> left, BinNode<E> right) {
        if (left == null && right == null) return null;
        else if (left == null) return right;
        else if (right == null) return left;
        else {
            boolean leftWin = random.nextInt(left.size + right.size) < left.size;
            if (leftWin) {
                left.bindRight(join(left.right, right));
                left.size += right.size;
                return left;
            }
            else {
                right.bindLeft(join(left, right.left));
                right.size += left.size;
                return right;
            }
        }
    }

    /**
     * If root != null, return the last non-null node in the searching path for key.
     * @param key
     * @return
     */
    private BinNode<E> find(BinNode<E> top, E element) {
        if (top == null) return null;

        BinNode<E> n = top;
        while (true) {
            int c = comparator.compare(element, n.element);
            if (c == 0) {
                return n;
            }
            BinNode<E> next = c < 0 ? n.left : n.right;
            if (next == null) return n;
            n = next;
        }
    }

    public boolean isValid() {
        return root == null || isValid(root);
    }

    private boolean isValid(BinNode<E> node) {
        if (node.left != null) {
            if (comparator.compare(node.element, node.left.element) <= 0) {
                return false;
            }
        }
        else if (node.right != null) {
            if (comparator.compare(node.element, node.right.element) >= 0) {
                return false;
            }
        }
        return node.size == 1 + (node.left != null ? node.left.size : 0) + (node.right != null ? node.right.size : 0);
    }

    @Override
    public String toString() {
        return TreePrinter.toString(root, BinNode::getChildren);
    }

}
